# -*- coding: utf-8 -*-
import jieba
import pandas as pd
import pickle
from pyexpat import features
import warnings
from settings import path

warnings.filterwarnings("ignore")
pd.set_option('display.max_columns', None)


def predict(data):
    with open(path.path_random_tree_model, 'rb') as f:
        model = pickle.load(f)
    with open(path.path_modol_random_tree_vectorizer, 'rb') as f:
        vectorizer = pickle.load(f)
    words = " ".join(jieba.lcut(data["text"])[:38])
    fearture = vectorizer.transform([words])
    y_pred = model.predict(fearture)
    id2class = {i: line.strip() for i, line in enumerate(open(path.path_class_txt, encoding='utf-8'))}
    data['pred_class'] = id2class[y_pred[0]]
    return data


# text = predict({'text': '天下万般兵刃，唯有过往伤人最深'})
# print(text)
